# -*- coding: utf-8 -*-
import time
import cv2
import numpy as np
from skimage.metrics import peak_signal_noise_ratio as psnr
from skimage.metrics import structural_similarity as ssim
import os
import sys

# Add algorithm module path
sys.path.append(os.path.join(os.path.dirname(__file__), 'algorithms'))

try:
    from lsb import LSB
    from dct import DCT
    from pvd import PVD
    print("[OK] Successfully imported all algorithm modules")
except ImportError as e:
    print(f"[ERROR] Error importing modules: {e}")
    print("Please check if algorithm files exist in 'algorithms' folder")
    sys.exit(1)

class PerformanceTester:
    def __init__(self):
        self.test_results = {}
        self.test_count = 0
        
    def create_test_image(self, size=(512, 512), filename="test_image.png"):
        """Create test image with rich texture"""  # 创建包含丰富纹理的测试图像
        try:
            # Create gradient image for better testing
            img = np.zeros((size[0], size[1], 3), dtype=np.uint8)
            for i in range(3):  # For each color channel
                for x in range(size[0]):
                    for y in range(size[1]):
                        img[x, y, i] = (x + y) % 256
            
            cv2.imwrite(filename, img)
            print(f"[OK] Created test image: {filename}")
            return filename
        except Exception as e:
            print(f"[ERROR] Error creating test image: {e}")
            return None
    
    def generate_test_data(self, length):
        """Generate test data with English and Chinese characters"""  # 生成包含英文和中文的测试数据
        try:
            import string
            import random
            
            # English part
            english_chars = string.ascii_letters + string.digits + " "
            english_part = ''.join(random.choices(english_chars, k=length//2))
            
            # Chinese part - use simple Chinese characters
            simple_chinese = ["Test", "Info", "Hide", "Algo", "Perf", "Eval", "System", "Module", "Func", "Data"]
            chinese_part = ''.join(random.choices(simple_chinese, k=length//10))
            
            # Combine
            test_data = english_part + chinese_part
            
            # Ensure exact length
            test_data = test_data[:length]
            
            print(f"[OK] Generated test data: {test_data[:30]}...")
            return test_data
        except Exception as e:
            print(f"[ERROR] Error generating test data: {e}")
            return "Test data " * (length // 10)  # Fallback data
    
    def calculate_image_quality(self, original_path, stego_path):
        """Calculate image quality metrics"""  # 计算图像质量指标
        try:
            original = cv2.imread(original_path)
            stego = cv2.imread(stego_path)
            
            if original is None:
                print(f"[ERROR] Cannot read original image: {original_path}")
                return 0, 0, 0
            if stego is None:
                print(f"[ERROR] Cannot read stego image: {stego_path}")
                return 0, 0, 0
            
            # Ensure same dimensions
            if original.shape != stego.shape:
                stego = cv2.resize(stego, (original.shape[1], original.shape[0]))
            
            # Calculate PSNR
            psnr_value = psnr(original, stego)
            
            # Calculate SSIM
            original_gray = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
            stego_gray = cv2.cvtColor(stego, cv2.COLOR_BGR2GRAY)
            ssim_value = ssim(original_gray, stego_gray, 
                            data_range=stego_gray.max() - stego_gray.min())
            
            # Calculate MSE
            mse_value = np.mean((original.astype(float) - stego.astype(float)) ** 2)
            
            print(f"[OK] Quality - PSNR: {psnr_value:.2f}dB, SSIM: {ssim_value:.4f}, MSE: {mse_value:.4f}")
            return psnr_value, ssim_value, mse_value
            
        except Exception as e:
            print(f"[ERROR] Error calculating quality metrics: {e}")
            return 0, 0, 0
    
    def test_single_algorithm(self, algorithm, algorithm_name, test_cases):
        """Test single algorithm performance"""  # 测试单个算法性能
        print("\n" + "="*60)
        print(f"Testing {algorithm_name} Algorithm")
        print("="*60)
        
        results = {}
        
        for case_name, test_data in test_cases.items():
            self.test_count += 1
            print(f"\n--- Test Case {self.test_count}: {case_name} ---")
            print(f"Data length: {len(test_data)} characters")
            print(f"Sample data: {test_data[:30]}")
            
            # Create test image
            test_image_path = self.create_test_image()
            if test_image_path is None:
                continue
                
            output_path = f"stego_{algorithm_name}_{case_name.replace('(', '').replace(')', '')}.png"
            
            # Embedding test
            print("Starting embedding test...")
            embed_start = time.time()
            try:
                embed_success = algorithm.embed(test_image_path, test_data, output_path)
                embed_time = time.time() - embed_start
                print(f"[OK] Embedding {'successful' if embed_success else 'failed'}")
            except Exception as e:
                print(f"[ERROR] Embedding error: {e}")
                embed_success = False
                embed_time = 0
            
            # Extraction test
            print("Starting extraction test...")
            extract_start = time.time()
            extracted_data = ""
            if embed_success:
                try:
                    extracted_data = algorithm.extract(output_path)
                    extract_time = time.time() - extract_start
                    print(f"[OK] Extraction completed")
                except Exception as e:
                    print(f"[ERROR] Extraction error: {e}")
                    extracted_data = ""
                    extract_time = 0
            else:
                extract_time = 0
            
            # Calculate image quality
            if embed_success and os.path.exists(output_path):
                psnr_val, ssim_val, mse_val = self.calculate_image_quality(test_image_path, output_path)
            else:
                psnr_val, ssim_val, mse_val = 0, 0, 0
            
            # Verify extraction result
            extraction_correct = False
            if embed_success and extracted_data:
                extraction_correct = (extracted_data == test_data)
                print(f"Extraction verification: {'PASS' if extraction_correct else 'FAIL'}")
                if not extraction_correct:
                    print(f"Original: {test_data[:30]}")
                    print(f"Extracted: {extracted_data[:30]}")
            
            # Store results
            results[case_name] = {
                'embed_time': embed_time,
                'extract_time': extract_time,
                'psnr': psnr_val,
                'ssim': ssim_val,
                'mse': mse_val,
                'embed_success': embed_success,
                'extraction_correct': extraction_correct,
                'original_data_length': len(test_data),
                'extracted_data_length': len(extracted_data) if extracted_data else 0
            }
            
            # Cleanup temporary files
            if os.path.exists(test_image_path):
                os.remove(test_image_path)
        
        return results
    
    def run_comprehensive_test(self):
        """Run comprehensive performance test"""  # 运行全面性能测试
        print("Starting Steganography Algorithm Performance Test")
        print("="*70)
        
        # Define test cases
        test_cases = {
            'Short Text (50 chars)': self.generate_test_data(50),      # 短文本
            'Medium Text (200 chars)': self.generate_test_data(200),   # 中等文本
            'Long Text (500 chars)': self.generate_test_data(500)      # 长文本
        }
        
        # Initialize algorithm instances
        algorithms = {
            'LSB': LSB(),
            'DCT': DCT(),
            'PVD': PVD()
        }
        
        # Run tests
        all_results = {}
        for algo_name, algorithm in algorithms.items():
            all_results[algo_name] = self.test_single_algorithm(algorithm, algo_name, test_cases)
        
        # Save results
        self.test_results = all_results
        
        # Generate report
        self.generate_comprehensive_report()
        
        return all_results
    
    def generate_comprehensive_report(self):
        """Generate comprehensive test report"""  # 生成全面测试报告
        print("\n" + "="*90)
        print("STEGANOGRAPHY ALGORITHMS PERFORMANCE TEST REPORT")
        print("="*90)
        
        # Performance comparison table
        print("\nPERFORMANCE COMPARISON TABLE:")
        print("-" * 110)
        header = f"{'Algorithm':<8} {'Test Case':<18} {'Embed(s)':<10} {'Extract(s)':<10} {'PSNR(dB)':<10} {'SSIM':<8} {'MSE':<10} {'Status':<8}"
        print(header)
        print("-" * 110)
        
        for algo_name, algo_results in self.test_results.items():
            for case_name, result in algo_results.items():
                status = "PASS" if result['extraction_correct'] else "FAIL"
                row = f"{algo_name:<8} {case_name:<18} {result['embed_time']:<10.4f} {result['extract_time']:<10.4f} " \
                      f"{result['psnr']:<10.2f} {result['ssim']:<8.4f} {result['mse']:<10.4f} {status:<8}"
                print(row)
        
        # Statistical analysis
        print("\nSTATISTICAL ANALYSIS:")
        print("-" * 60)
        
        for algo_name, algo_results in self.test_results.items():
            embed_times = [r['embed_time'] for r in algo_results.values() if r['embed_success']]
            extract_times = [r['extract_time'] for r in algo_results.values() if r['embed_success']]
            psnrs = [r['psnr'] for r in algo_results.values() if r['embed_success']]
            ssims = [r['ssim'] for r in algo_results.values() if r['embed_success']]
            success_count = sum([1 for r in algo_results.values() if r['extraction_correct']])
            success_rate = success_count / len(algo_results) * 100 if algo_results else 0
            
            if embed_times:  # Check if list is not empty
                print(f"\n{algo_name} Algorithm:")
                print(f"  Average Embed Time: {np.mean(embed_times):.4f}s")
                print(f"  Average Extract Time: {np.mean(extract_times):.4f}s")
                print(f"  Average PSNR: {np.mean(psnrs):.2f} dB")
                print(f"  Average SSIM: {np.mean(ssims):.4f}")
                print(f"  Success Rate: {success_rate:.1f}% ({success_count}/{len(algo_results)})")
            else:
                print(f"\n{algo_name}: No successful tests")
        
        # Performance ranking
        print("\nPERFORMANCE RANKING:")
        print("-" * 40)
        
        # Calculate overall scores
        algo_scores = {}
        for algo_name, algo_results in self.test_results.items():
            if algo_results:
                success_rate = sum([1 for r in algo_results.values() if r['extraction_correct']]) / len(algo_results) * 100
                avg_psnr = np.mean([r['psnr'] for r in algo_results.values() if r['embed_success']]) if any(r['embed_success'] for r in algo_results.values()) else 0
                avg_speed = 1 / (np.mean([r['embed_time'] for r in algo_results.values() if r['embed_success']]) + 0.001) if any(r['embed_success'] for r in algo_results.values()) else 0
                
                # Combined score (higher is better)
                score = success_rate * 0.5 + avg_psnr * 0.3 + avg_speed * 0.2
                algo_scores[algo_name] = score
        
        # Sort by score
        sorted_algos = sorted(algo_scores.items(), key=lambda x: x[1], reverse=True)
        
        for i, (algo_name, score) in enumerate(sorted_algos, 1):
            print(f"{i}. {algo_name}: {score:.2f} points")
        
        # Recommendations
        print("\nRECOMMENDATIONS:")
        print("-" * 40)
        print("1. LSB: Best for speed and capacity")        # 最适合速度和容量
        print("2. DCT: Best for robustness")                # 最适合鲁棒性
        print("3. PVD: Balanced performance")               # 平衡性能
        print("4. Choose based on application requirements") # 根据应用需求选择

def main():
    """Main function"""  # 主函数
    try:
        tester = PerformanceTester()
        results = tester.run_comprehensive_test()
        
        print("\n" + "="*70)
        print("PERFORMANCE TEST COMPLETED")
        print("="*70)
        
        return results
        
    except Exception as e:
        print(f"\n[ERROR] Performance test failed: {e}")
        print("Please check algorithm implementations")
        return None

if __name__ == "__main__":
    # Run the test
    print("Initializing Performance Test...")
    results = main()
    
    # Keep console open
    if results:
        input("\nPress Enter to exit...")
    else:
        input("\nTest completed with errors. Press Enter to exit...")